## Private Quantile estimation

estimateQuantile <- function(target_quantile, rng, epsilon, data) {
  max <- rng[2]
  min <- rng[1]
  num_clients <- length(data)

  sorted_data <- sort(data)
  extended_data <- c(min, sorted_data, max)
  y <- c()
  for (i in 1:(length(extended_data) - 1)) {
    y[i] <- log((extended_data[i + 1] - extended_data[i])) - epsilon * abs(i - target_quantile * num_clients)
  }
  random_integer <- gumbelSample(y)
  quantile_estimate <- runif(1, extended_data[random_integer], extended_data[random_integer + 1])
  return(quantile_estimate)
}

gumbelSample <- function(x) {
  z <- evd::rgumbel(length(x), loc = 0, scale = 1)
  w <- x + z
  return(which.max(w))
}


widenedMean <- function(rng, epsilon, data) {
  P <- length(data)
  rad <- P^(1 / 3 + 1 / 10)
  a_hat <- estimateQuantile(target_quantile = 0.25, rng = c(0, 20), epsilon = epsilon / 4, data = data)
  b_hat <- estimateQuantile(target_quantile = 0.75, rng = c(0, 20), epsilon = epsilon / 4, data = data)
  mu_crude <- (a_hat + b_hat) / 2
  iqr_crude <- abs(b_hat - a_hat)
  u <- mu_crude + 4 * rad * iqr_crude
  l <- mu_crude - 4 * rad * iqr_crude
  censored_data <- ifelse(data > u, u, data)
  censored_data <- ifelse(censored_data < l, l, censored_data)
  mu_hat <- mean(censored_data)
  return(mu_hat + rmutil::rlaplace(n = 1, m = 0, s = abs(u - l) / ((epsilon / 2) * P)))
}

## Run simulations
n <- 10000
P <- sqrt(n) 

simulationFn <- function(epsilon_mean, epsilon_quantile, target_quantile) {
  
  dat <- rpois(n, lambda = 5)

  dat_list <- splitData(N = n, n = n / P, P = P)

  partition_estimate <- c()

  for (i in 1:length(dat_list)) {
    partition_estimate[i] <- mean(dat[dat_list[[i]]])
  }


  widened_mean <- widenedMean(rng = c(0, 10), epsilon = epsilon_mean + epsilon_quantile, data = partition_estimate)

  lambda <- estimateQuantile(target_quantile = target_quantile, rng = c(0, 10), epsilon = epsilon_quantile, data = partition_estimate)


  cen <- censorParam(partition_estimate, lambda = lambda, delta = 1 / n, epsilon = epsilon_mean, epsilon_alpha = epsilon_mean)
  theta_hat_dp <- cen[[1]]
  alpha_noise <- cen[[4]]


  beta_tilde <- biasAdjustment(
    theta_dp = theta_hat_dp, alpha = alpha_noise,
    lambda = lambda, upper = TRUE, two_sided = F
  )

  return(c(
    theta_tilde = beta_tilde$theta_tilde,
    theta_hat = theta_hat_dp,
    widened_mean = widened_mean,
    epsilon_mean = epsilon_mean,
    epsilon_quantile = epsilon_quantile
  ))
}

set.seed(02138)

# Run sims
sims_1 <- t(replicate(1000, simulationFn(epsilon_mean = 0.5, epsilon_quantile = 0.5, target_quantile = 0.6)))
sims_2 <- t(replicate(1000, simulationFn(epsilon_mean = 0.75, epsilon_quantile = 0.75, target_quantile = 0.6)))
sims_3 <- t(replicate(1000, simulationFn(epsilon_mean = 1, epsilon_quantile = 1, target_quantile = 0.6)))
sims_4 <- t(replicate(1000, simulationFn(epsilon_mean = 1.25, epsilon_quantile = 1.25, target_quantile = 0.6)))
sims_5 <- t(replicate(1000, simulationFn(epsilon_mean = 1.5, epsilon_quantile = 1.5, target_quantile = 0.6)))

sims_df <- rbind(sims_1, sims_2, sims_3, sims_4, sims_5) %>% as.data.frame()

sims_summary <- sims_df %>%
  group_by(epsilon_mean) %>%
  summarize(
    rmse_tilde = mean(theta_tilde - 5)^2 + var(theta_tilde),
    rmse_hat = mean(theta_hat - 5)^2 + var(theta_hat),
    rmse_widened = mean(widened_mean - 5)^2 + var(widened_mean),
    mean(theta_tilde),
    mean(theta_hat),
    mean(widened_mean)
  )



plot <- sims_summary %>%
  ggplot() +
  geom_point(aes(x = epsilon_mean, y = log(rmse_tilde), col = "Bias adjusted")) +
  geom_line(aes(x = epsilon_mean, y = log(rmse_tilde), col = "Bias adjusted")) +
  # geom_point(aes(x = epsilon_mean*2, y = log(rmse_hat), col = 'theta_hat')) +
  # geom_line(aes(x = epsilon_mean*2, y = log(rmse_hat), col = 'theta_hat'))  +
  geom_point(aes(x = epsilon_mean, y = log(rmse_widened), col = "Smith (2011)")) +
  geom_line(aes(x = epsilon_mean, y = log(rmse_widened), col = "Smith (2011)")) +
  theme_bw() +
  theme(
    legend.position = c(0.9, 0.9),
    legend.justification = c(0.9, 0.9),
    panel.grid.major = element_blank(),
    legend.box.background = element_rect(colour = "black")
  ) +
  labs(x = expression(epsilon), y = "RMSE (logged)", col = "")

plot
ggsave("figA1", plot, height = 4, width = 5)
